#!/bin/bash

# 通用rsync下载脚本
# 用法: ./download.sh [远程路径] [本地路径] [SSH主机] [SSH端口] [SSH用户]
# 默认: ~/project -> 当前目录（下载outputs和mlflow相关文件夹）

set -e

# 默认参数
REMOTE_PATH="${1:-~/project}"
LOCAL_PATH="${2:-.}"
SSH_HOST="${3:-northwest1.gpugeek.com}"
SSH_PORT="${4:-56740}"
SSH_USER="${5:-root}"

# 颜色输出
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

echo -e "${YELLOW}=== Rsync 下载脚本 ===${NC}"
echo -e "远程路径: ${GREEN}$REMOTE_PATH${NC}"
echo -e "本地路径: ${GREEN}$LOCAL_PATH${NC}"
echo -e "SSH连接: ${GREEN}$SSH_USER@$SSH_HOST:$SSH_PORT${NC}"

# 添加连接测试功能
echo -e "${YELLOW}测试SSH连接...${NC}"
if ! ssh -p $SSH_PORT -o ConnectTimeout=10 $SSH_USER@$SSH_HOST "echo 'SSH连接成功'" 2>/dev/null; then
    echo -e "${RED}❌ SSH连接失败！请检查主机、端口、用户名是否正确${NC}"
    exit 1
fi
echo -e "${GREEN}✅ SSH连接测试成功${NC}"

# 下载outputs和mlflow相关文件夹
echo -e "${YELLOW}开始下载outputs和mlflow相关文件夹...${NC}"

rsync -avz --progress \
    -e "ssh -p $SSH_PORT" \
    --include="output/" \
    --include="outputs/" \
    --include="outputs/***" \
    --include="mlruns/" \
    --include="mlruns/***" \
    --exclude="*" \
    "$SSH_USER@$SSH_HOST:$REMOTE_PATH/" "$LOCAL_PATH/"

echo -e "${GREEN}✅ 下载完成！${NC}"
echo -e "${YELLOW}连接服务器: ssh -p $SSH_PORT $SSH_USER@$SSH_HOST${NC}"
